/** @file admm_coord_descent_sub_problem_solver.h
 *
 *  Coordinate descent based ADMM sub problem solver.
 *
 *  @author Dongryeol Lee (dongryel@cc.gatech.edu)
 */

#ifndef MLPACK_ADMM_ADMM_COORD_DESCENT_SUB_PROBLEM_SOLVER_H
#define MLPACK_ADMM_ADMM_COORD_DESCENT_SUB_PROBLEM_SOLVER_H

#include "core/data/matrix_dataset.h"
#include "core/math/math.h"
#include "mlpack/admm/admm_result.h"
#include "mlpack/admm/admm_sub_problem_solver.h"

namespace mlpack {
namespace admm {
class ADMMCoordDescentSubProblemSolver :
  public mlpack::admm::ADMMSubProblemSolver {
  public:

    void Solve(
      const mlpack::admm::ADMMArguments &arguments_in,
      mlpack::admm::ADMMResult *result_out) const {

      const double convergence_tolerance = 1e-2;

      // The training dataset with the points and their labels.
      core::data::MatrixDataset *reference_table =
        const_cast <
        mlpack::admm::ADMMArguments & >(arguments_in).reference_table_;

      // Initialize the current local primal variables.
      arma::vec current_local_primal_variables =
        result_out->local_adjusted_lagrangian_variables_ +
        result_out->synchronized_lagrangian_variables_ / arguments_in.admm_rho_;

      // Difference between $v_k^t$ and $\omega^t$.
      arma::vec diff_lagrangian_variables =
        result_out->local_adjusted_lagrangian_variables_ +
        result_out->synchronized_lagrangian_variables_ / arguments_in.admm_rho_;

      // Temporary variables.
      int num_converged_coordinates = 0;
      arma::vec diff_local_lagrangian_variables;

      // Repeat until convergence.
      int num_iterations = 0;
      while(num_converged_coordinates < reference_table->num_points()) {

        // Increment the number of iterations.
        num_iterations++;

        // Reset the number of converged coordinates.
        num_converged_coordinates = 0;

        // Cyclic-loop each coordinate.
        for(int i = 0; i < reference_table->num_points(); i++) {

          // Get the point and its label (assumed to be the last
          // element).
          const core::data::alias_vec *targets = NULL;
          const core::data::alias_vec &point =
            reference_table->get_instance(i, &targets);
          int label = static_cast<int>((*targets)[0]);

          // Determine whether the partial derivative along this
          // coordinate is zero. If so, increment the number of
          // converged coordinates.
          diff_local_lagrangian_variables =
            current_local_primal_variables - diff_lagrangian_variables;
          double projected_partial_derivative =
            label * arma::dot(
              point, diff_local_lagrangian_variables) +
            result_out->local_dual_variables_[i] /
            (2 * arguments_in.svm_c_parameter_ * arguments_in.admm_rho_)   -
            arguments_in.admm_rho_inverse_;

          if(result_out->local_dual_variables_[i] <= convergence_tolerance) {
            projected_partial_derivative =
              std::min(0.0, projected_partial_derivative);
          }

          // If converged, then skip the current coordinate.
          if(fabs(projected_partial_derivative) <= convergence_tolerance) {
            num_converged_coordinates++;
            continue;
          }

          // Move along the dimension corresponding to this point.
          double numerator =
            arma::dot(diff_local_lagrangian_variables, point) * label -
            arguments_in.admm_rho_inverse_;

          // A common factor that needs to be divided in coordinate step.
          double denominator =
            arguments_in.squared_norm_reference_point(i) +
            1.0 / (
              2 * arguments_in.svm_c_parameter_ * arguments_in.admm_rho_);

          double old_dual_coord_value = result_out->local_dual_variables_[i];
          result_out->local_dual_variables_[i] =
            (arguments_in.squared_norm_reference_point(i) *
             old_dual_coord_value - numerator) / denominator;

          // Thresholding.
          result_out->local_dual_variables_[i] =
            std::max(0.0, result_out->local_dual_variables_[i]);

          // Update the Lagrangian variable here based on the change in
          // the current dual variable coordinate.
          core::math::scaled_add(
            label *
            (result_out->local_dual_variables_[i] - old_dual_coord_value),
            point, &current_local_primal_variables);
        }
      }
    }
};
}
}

#endif
